/****************************************************************************
 * Copyright (C) 2009-2010 SciTouch LLC
 * 
 * This file is part of Indigo toolkit.
 * 
 * This file may be distributed and/or modified under the terms of the
 * GNU General Public License version 3 as published by the Free Software
 * Foundation and appearing in the file LICENSE.GPL included in the
 * packaging of this file.
 * 
 * This file is provided AS IS with NO WARRANTY OF ANY KIND, INCLUDING THE
 * WARRANTY OF DESIGN, MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE.
 ***************************************************************************/

#include "reaction/reaction.h"
#include "molecule/molecule_arom_match.h"
#include "molecule/molecule_3d_constraints.h"
#include "molecule/molecule_substructure_matcher.h"
#include "reaction/reaction_substructure_matcher.h"

ReactionSubstructureMatcher::ReactionSubstructureMatcher (const Reaction &target) :
_target(target),
TL_CP_GET(_matchers),
TL_CP_GET(_amm_to_second_side_1),
TL_CP_GET(_amm_to_second_side_2),
TL_CP_GET(_molecule_core_1),
TL_CP_GET(_molecule_core_2),
TL_CP_GET(_amm_core_first_side),
TL_CP_GET(_query_h_counters),
TL_CP_GET(_target_h_counters)
{
   use_aromaticity_matcher = true;
   highlighting = 0;

   _query_nei_counters = 0;
   _target_nei_counters = 0;

   _query = 0;

   _target_h_counters.resize(_target.end());

   int i;

   for (i = _target.begin(); i < _target.end(); i = _target.next(i))
      MoleculeSubstructureMatcher::calcTargetHCounters(_target.getMolecule(i), _target_h_counters[i]);

   _matchers.clear();
   _matchers.push(*this);
}

void ReactionSubstructureMatcher::setQuery (const Reaction &query)
{
   _query = &query;

   _query_h_counters.resize(query.end());

   int i;

   for (i = query.begin(); i < query.end(); i = query.next(i))
      MoleculeSubstructureMatcher::calcQueryHCounters(query.getMolecule(i), _query_h_counters[i], 0);
}

void ReactionSubstructureMatcher::setNeiCounters (const ReactionAtomNeighbourhoodCounters *query_counters, 
                                                  const ReactionAtomNeighbourhoodCounters *target_counters)
{
   _query_nei_counters = query_counters;
   _target_nei_counters = target_counters;
}

bool ReactionSubstructureMatcher::find ()
{
   if (_query == 0)
      throw Error("no query");

   if (highlighting != 0)
      highlighting->init(_target);

   if (_query->reactantsCount() > _target.reactantsCount() || _query->productsCount() > _target.productsCount())
      return false;

   if (_query->reactantsCount() * _target.reactantsCount() < _query->productsCount() * _target.productsCount())
      _first_side = Reaction::REACTANT, _second_side = Reaction::PRODUCT;
   else 
      _first_side = Reaction::PRODUCT, _second_side = Reaction::REACTANT;

   _initMap(*_query, _amm_to_second_side_1);
   _initMap(_target, _amm_to_second_side_2);

   _molecule_core_1.resize(_query->end());
   _molecule_core_1.fffill();
   _molecule_core_2.resize(_target.end());
   _molecule_core_2.fffill();
   _amm_core_first_side.clear();

   while (1)
   {
      int command = _matchers.top().nextPair();

      if (command == _CONTINUE)
         continue;
      
      if (command == _RETURN)
      {
         if (_checkAAM())
         {
            _highlight();
            while (_matchers.size() > 1)
               _matchers.pop();
            return true;
         }
         command = _NO_WAY;
      } else if (command != _NO_WAY)
      {
         int mol1 = _matchers.top()._current_molecule_1;
         int mol2 = _matchers.top()._current_molecule_2;
         Array<int> &core1 = _matchers.top()._current_core_1;
         Array<int> &core2 = _matchers.top()._current_core_2;
         int mode = _matchers.top().getMode();

         _matchers.reserve(_matchers.size() + 1);
         _matchers.push(_matchers.top());
         _matchers.top().setMode(command);
         if (!_matchers.top().addPair(mol1, mol2, core1, core2, mode == _FIRST_SIDE))
            _matchers.pop();
      }

      if (command == _NO_WAY)
      {
         if (_matchers.size() > 1)
         {
            _matchers.top().restore();
            _matchers.pop();
         }
         else
            return false;
      }
   }
}

// Init data for reaction substructure search
void ReactionSubstructureMatcher::_initMap (const Reaction &reaction, RedBlackMap<int, int> &aam_map)
{
   int aam_number;
   int i, j;
   int *val;

   aam_map.clear();

   // collect aam-to-molecule index mapping for reaction second side
   for (i = reaction.sideBegin(_second_side); i < reaction.sideEnd(); i = reaction.sideNext(_second_side, i))
   {
      const Molecule &i_mol = reaction.getMolecule(i);

      for (j = i_mol.vertexBegin(); j < i_mol.vertexEnd(); j = i_mol.vertexNext(j))
         if ((aam_number = reaction.getAAM(i, j)) != 0)
         {
            if ((val = aam_map.at2(aam_number)) == 0)
               aam_map.insert(aam_number, i);
            else if (*val < 0)
               (*val)--;
            else
               (*val) = -1;
         }
   }
}

// Check correct AAM relationship between query and target reaction
int ReactionSubstructureMatcher::_checkAAM () const
{
   int *aam, aam1, aam2;
   int i, j;

   for (i = 1; i < _matchers.size() - 1; i++)
   {
      if (_matchers[i].getMode() == _FIRST_SIDE)
         continue;

      const Molecule &mol1 = _query->getMolecule(_matchers[i]._current_molecule_1);

      for (j = mol1.vertexBegin(); j < mol1.vertexEnd(); j = mol1.vertexNext(j))
      {
         int k = _matchers[i]._current_core_1[j];

         if (k < 0)
            continue;

         aam1 = _query->getAAM(_matchers[i]._current_molecule_1, j);
         aam2 = _target.getAAM(_matchers[i]._current_molecule_2, k);

         if (aam1 > 0 && aam2 > 0)
            if ((aam = _amm_core_first_side.at2(aam1)) != 0 && *aam != aam2)
               return false;
      }
   }

   return true;
}

void ReactionSubstructureMatcher::_highlight ()
{
   if (highlighting == 0)
      return;

   int i;

   for (i = 0; i < _matchers.size() - 1; i++)
      highlighting->getGraphHighlighting(_matchers[i]._current_molecule_2).onSubgraph(_query->getMolecule(_matchers[i]._current_molecule_1), _matchers[i]._current_core_1.ptr());
}

ReactionSubstructureMatcher::_Matcher::_Matcher (ReactionSubstructureMatcher &context) :
TL_CP_GET(_current_core_1),
TL_CP_GET(_current_core_2),
_context(context),
TL_CP_GET(_mapped_aams)
{
   _mode = _FIRST_SIDE;
   _selected_molecule_1 = -1;
   _selected_molecule_2 = -1;
   _current_molecule_1 = -1;
   _current_molecule_2 = -1;
   _mapped_aams.clear();
}

ReactionSubstructureMatcher::_Matcher::_Matcher (const ReactionSubstructureMatcher::_Matcher &other) :
TL_CP_GET(_current_core_1),
TL_CP_GET(_current_core_2),
_context(other._context),
TL_CP_GET(_mapped_aams)
{
   _current_molecule_1 = -1;
   _current_molecule_2 = -1;
   _mapped_aams.clear();
}

int ReactionSubstructureMatcher::_Matcher::_nextPair ()
{
   int side;
   int i;

   if (_mode == _FIRST_SIDE)
      side = _context._first_side;
   else // _SECOND_SIDE_REST
      side = _context._second_side;

   if (_enumerator.get() == 0 || !_enumerator->processNext())
   {
      do
      {
         if (_current_molecule_1 == -1)
         {
            for (_current_molecule_1 = _context._query->sideBegin(side); 
               _current_molecule_1 < _context._query->sideEnd(); 
               _current_molecule_1 = _context._query->sideNext(side, _current_molecule_1))
               if (_context._molecule_core_1[_current_molecule_1] < 0)
                  break;
            if (_current_molecule_1 == _context._query->sideEnd())
            {
               if (_mode == _FIRST_SIDE)
               {
                  _mode = _SECOND_SIDE_REST;
                  _current_molecule_1 = -1;
                  return _nextPair();
               }

               return _RETURN;
            }
         }

         if (_current_molecule_2 == -1)
            _current_molecule_2 = _context._target.sideBegin(side);
         else
            _current_molecule_2 = _context._target.sideNext(side, _current_molecule_2);

         for ( ; _current_molecule_2 < _context._target.sideEnd(); _current_molecule_2 = _context._target.sideNext(side, _current_molecule_2))
            if (_context._molecule_core_2[_current_molecule_2] < 0)
               break;

         if (_current_molecule_2 == _context._target.sideEnd())
            return _NO_WAY;

         _enumerator.free();

         const Molecule &mol_1 = _context._query->getMolecule(_current_molecule_1);
         const Molecule &mol_2 = _context._target.getMolecule(_current_molecule_2);

         _enumerator.create(mol_2);

         // init embedding enumerator context
         _enumerator->cb_match_edge = _matchBonds;
         _enumerator->cb_match_vertex = _matchAtoms;
         _enumerator->cb_edge_add = _addBond;
         _enumerator->cb_vertex_add = _addAtom;
         _enumerator->cb_vertex_remove = _removeAtom;
         _enumerator->cb_embedding = _embedding;

         // mark hydrogens to ignore
         QS_DEF(Array<int>, ignored);

         ignored.clear_resize(mol_1.vertexEnd());

         MoleculeSubstructureMatcher::markIgnoredHydrogens(mol_1, ignored.ptr(), 0, 1);

         if (_context.use_aromaticity_matcher)
            _am.reset(new AromaticityMatcher(mol_1, mol_2));

         _enumerator->userdata = this;
         _enumerator->setSubgraph(mol_1);

         for (i = mol_1.vertexBegin(); i != mol_1.vertexEnd(); i = mol_1.vertexNext(i))
            if (ignored[i])
               _enumerator->ignoreSubgraphVertex(i);

         _enumerator->processStart();

      } while (!_enumerator->processNext());
   }

   return _mode == _FIRST_SIDE ? _SECOND_SIDE : _SECOND_SIDE_REST;
}

int ReactionSubstructureMatcher::_Matcher::nextPair ()
{
   if (_mode != _SECOND_SIDE)
   {
      int next = _nextPair();

      if (next != _SECOND_SIDE)
         return next;

      // Switch to _SECOND_SIDE
      const Molecule &mol_1 = _context._query->getMolecule(_current_molecule_1);

      int first_aam_1 = 0;
      int first_aam_2 = 0;

      int i;

      for (i = mol_1.vertexBegin(); i < mol_1.vertexEnd(); i = mol_1.vertexNext(i))
         if (_current_core_1[i] >= 0)
         {
            first_aam_1 = _context._query->getAAM(_current_molecule_1, i);
            first_aam_2 = _context._target.getAAM(_current_molecule_2, _current_core_1[i]);
            break;
         }


      if (first_aam_1 > 0 && first_aam_2 > 0)
      {
         // Check the other side if need
         int mol_1_idx_ss = _context._amm_to_second_side_1.at(first_aam_1);
         int mol_2_idx_ss = _context._amm_to_second_side_2.at(first_aam_2);

         if ((mol_1_idx_ss < 0 && mol_1_idx_ss < mol_2_idx_ss))
            return _CONTINUE; // subreactions equal AAM-numbers more than superreaction

         if (mol_2_idx_ss < 0)
            return _FIRST_SIDE; // check this molecules in the completion phase

         if (_context._molecule_core_1[mol_1_idx_ss] >= 0)
         {
            if (_context._molecule_core_1[mol_1_idx_ss] != mol_2_idx_ss)
               return _CONTINUE;

            int first_idx_1_ss = _context._query->findAtomByAAM(mol_1_idx_ss, first_aam_1);
            int first_idx_2_ss = _context._target.findAtomByAAM(mol_2_idx_ss, first_aam_2);
            int i;

            for (i = 0; i < _context._matchers.size(); i++)
               if (_context._matchers[i]._current_molecule_1 == mol_1_idx_ss)
               {
                  if (_context._matchers[i]._current_core_1[first_idx_1_ss] != first_idx_2_ss)
                     return _CONTINUE;
                  return _FIRST_SIDE;
               }
         }

         return _SECOND_SIDE;
      }

      return _FIRST_SIDE;
   }

   // _SECOND_SIDE
   if (_enumerator.get() == 0)
   {
      int i;

      const Molecule &src_mol_1 = _context._query->getMolecule(_selected_molecule_1);
      const Molecule &src_mol_2 = _context._target.getMolecule(_selected_molecule_2);

      int src_aam_1 = _context._query->findAamNumber(&src_mol_1, src_mol_1.vertexBegin());
      int src_aam_2 = _context._target.findAamNumber(&src_mol_2, _context._matchers[_context._matchers.size() - 2]._current_core_1[src_mol_1.vertexBegin()]);

      const Molecule &mol_1 = _context._query->getMolecule(_current_molecule_1);
      const Molecule &mol_2 = _context._target.getMolecule(_current_molecule_2);

      int first_idx_1 = _context._query->findAtomByAAM(_current_molecule_1, src_aam_1);
      int first_idx_2 = _context._target.findAtomByAAM(_current_molecule_2, src_aam_2);

      // init embedding enumerator context
      _enumerator.create(mol_2);

      _enumerator->cb_match_edge = _matchBonds;
      _enumerator->cb_match_vertex = _matchAtoms;
      _enumerator->cb_edge_add = _addBond;
      _enumerator->cb_vertex_add = _addAtom;
      _enumerator->cb_vertex_remove = _removeAtom;
      _enumerator->cb_embedding = _embedding;

      // mark hydrogens to ignore
      QS_DEF(Array<int>, ignored);

      ignored.clear_resize(mol_1.vertexEnd());

      MoleculeSubstructureMatcher::markIgnoredHydrogens(mol_1, ignored.ptr(), 0, 1);

      if (_context.use_aromaticity_matcher)
         _am.reset(new AromaticityMatcher(mol_1, mol_2));

      _enumerator->userdata = this;
      _enumerator->setSubgraph(mol_1);

      for (i = mol_1.vertexBegin(); i != mol_1.vertexEnd(); i = mol_1.vertexNext(i))
         if (ignored[i])
            _enumerator->ignoreSubgraphVertex(i);

      if (!_enumerator->fix(first_idx_1, first_idx_2))
         return _NO_WAY;

      _enumerator->processStart();
   }

   if (!_enumerator->processNext())
      return _NO_WAY;

   return _FIRST_SIDE;
}

bool ReactionSubstructureMatcher::_Matcher::addPair (int mol1_idx, int mol2_idx, const Array<int> &core1, const Array<int> &core2, bool from_first_side)
{
   _selected_molecule_1 = mol1_idx;
   _selected_molecule_2 = mol2_idx;

   _mapped_aams.clear();

   const Molecule &mol1 = _context._query->getMolecule(mol1_idx);

   if (from_first_side)
   {
      int i;

      for (i = mol1.vertexBegin(); i < mol1.vertexEnd(); i = mol1.vertexNext(i))
         if (core1[i] >= 0)
         {
            int aam1 = _context._query->getAAM(mol1_idx, i);
            int aam2 = _context._target.getAAM(mol2_idx, core1[i]);
            int *aam;

            if (aam1 > 0 && aam2 > 0)
            {
               if ((aam = _context._amm_core_first_side.at2(aam1)) == 0)
               {
                  _context._amm_core_first_side.insert(aam1, aam2);
                  _mapped_aams.push(aam1);
               } else if (*aam != aam2)
               {
                  while (_mapped_aams.size() > 0)
                     _context._amm_core_first_side.remove(_mapped_aams.pop());
                  return false;
               }
            }
         }
   }

   if (_mode == _SECOND_SIDE)
   {
      int first_aam_1 = 0;
      int first_aam_2 = 0;

      int i;

      for (i = mol1.vertexBegin(); i < mol1.vertexEnd(); i = mol1.vertexNext(i))
         if (core1[i] >= 0)
         {
            first_aam_1 = _context._query->getAAM(mol1_idx, i);
            first_aam_2 = _context._target.getAAM(mol2_idx, core1[i]);
            break;
         }

      _current_molecule_1 = _context._amm_to_second_side_1.at(first_aam_1);
      _current_molecule_2 = _context._amm_to_second_side_2.at(first_aam_2);
   }

   _context._molecule_core_1[mol1_idx] = mol2_idx;
   _context._molecule_core_2[mol2_idx] = mol1_idx;

   return true;
}

void ReactionSubstructureMatcher::_Matcher::restore()
{
   _context._molecule_core_1[_selected_molecule_1] = -1;
   _context._molecule_core_2[_selected_molecule_2] = -1;

   while (_mapped_aams.size() > 0)
      _context._amm_core_first_side.remove(_mapped_aams.pop());
}

int ReactionSubstructureMatcher::_Matcher::_embedding (const Graph &subgraph, const Graph &supergraph, int *core_sub, int *core_super, void *userdata)
{
   ReactionSubstructureMatcher::_Matcher &self = *(ReactionSubstructureMatcher::_Matcher *)userdata;

   const Molecule &query = (const Molecule &)subgraph;
   const Molecule &target = (const Molecule &)supergraph;

   if (!MoleculeStereocenters::checkSub(query.getStereocenters(), target.getStereocenters(), core_sub, false))
      return 1;

   if (!MoleculeCisTrans::checkSub(query, target, core_sub))
      return 1;

   self._current_core_1.copy(core_sub, subgraph.vertexEnd());
   self._current_core_2.copy(core_super, supergraph.vertexEnd());

   return 0;
}

bool ReactionSubstructureMatcher::_Matcher::_matchAtoms (const Graph &subgraph, const Graph &supergraph, const int *core_sub, int sub_idx, int super_idx, void *userdata)
{
   if (!MoleculeSubstructureMatcher::matchAtoms(subgraph, supergraph, core_sub, sub_idx, super_idx))
      return false;

   ReactionSubstructureMatcher::_Matcher *self = (ReactionSubstructureMatcher::_Matcher *)userdata;

   const Molecule &sub_mol = static_cast<const Molecule&>(subgraph);
   const Molecule &super_mol = static_cast<const Molecule&>(supergraph);

   if (self->_mode == _SECOND_SIDE)
   {
      int *aam, aam1, aam2;

      aam1 = self->_context._query->getAAM(self->_current_molecule_1, sub_idx);
      if (aam1 != 0)
      {
         aam2 = self->_context._target.getAAM(self->_current_molecule_2, super_idx);
         if (aam2 != 0)
            if ((aam = self->_context._amm_core_first_side.at2(aam1)) != 0 && *aam != aam2)
               return false;
      }
   }

   const Array<int> &query_h_counters = self->_context._query_h_counters[self->_current_molecule_1];
   const Array<int> &target_h_counters = self->_context._target_h_counters[self->_current_molecule_2];

   if (sub_idx < query_h_counters.size() && super_idx < target_h_counters.size())
      if (query_h_counters[sub_idx] > target_h_counters[super_idx])
         return false;

   const QueryAtom &sub_query_atom = sub_mol.getQueryAtom(sub_idx);

   if (sub_query_atom.hydrogens == 1) // means H0
   {
      if (target_h_counters[super_idx] > query_h_counters[sub_idx])
         return false;
   }
   else if (sub_query_atom.hydrogens > 1) // means Hn
   {
      if (target_h_counters[super_idx] < (sub_query_atom.hydrogens - 1))
         return false;
   }

   if (self->_context._query_nei_counters != 0 && self->_context._target_nei_counters != 0)
   {
      const MoleculeAtomNeighbourhoodCounters &mol_count1 = self->_context._query_nei_counters->getCounters(self->_current_molecule_1);
      const MoleculeAtomNeighbourhoodCounters &mol_count2 = self->_context._target_nei_counters->getCounters(self->_current_molecule_2);

      if (!mol_count1.testSubstructure(mol_count2, sub_idx, super_idx, true))
         return false;
   }

   int sub_atom_inv = self->_context._query->getInversion(self->_current_molecule_1, sub_idx);
   int super_atom_inv = self->_context._target.getInversion(self->_current_molecule_2, super_idx);

   if (sub_atom_inv != STEREO_UNMARKED && sub_atom_inv != super_atom_inv)
      return false;

   if (self->_context._query->getExactChange(self->_current_molecule_1, sub_idx) == 1)
   {
      const Vertex &can_vertex = sub_mol.getVertex(sub_idx);
      int ch_flag;
      int bonds_changes[RC_TOTAL] = {0};
      int i;

      for (i = can_vertex.neiBegin(); i != can_vertex.neiEnd(); i = can_vertex.neiNext(i))
      {
         ch_flag = self->_context._query->getReactingCenter(self->_current_molecule_1, can_vertex.neiEdge(i));

         if (ch_flag == RC_NOT_CENTER)
            ch_flag = RC_UNCHANGED;

         if (ch_flag > 0)
            bonds_changes[ch_flag]++;
      }

      const Vertex &pat_vertex = super_mol.getVertex(super_idx);

      for (i = pat_vertex.neiBegin(); i != pat_vertex.neiEnd(); i = pat_vertex.neiNext(i))
      {

         ch_flag = self->_context._target.getReactingCenter(self->_current_molecule_2, pat_vertex.neiEdge(i));
         if (ch_flag > 0)
            bonds_changes[ch_flag]--;
      }

      int n_centers = bonds_changes[RC_CENTER];
      bonds_changes[RC_CENTER] = 0;

      if (n_centers < 0)
         return false;

      for (int i = 0; i < NELEM(bonds_changes); i++)
      {
         if ((ch_flag = bonds_changes[i]) > 0)
            return false;
         else if (ch_flag < 0)
         {
            if ((n_centers += ch_flag) < 0)
               return false;
         }
      }

      if (n_centers != 0)
         return false;
   }

   return true;
}

bool ReactionSubstructureMatcher::_Matcher::_matchBonds (const Graph &subgraph, const Graph &supergraph, int sub_idx, int super_idx, void *userdata)
{
   ReactionSubstructureMatcher::_Matcher *self = (ReactionSubstructureMatcher::_Matcher *)userdata;

   if (!MoleculeSubstructureMatcher::matchBonds(subgraph, supergraph, sub_idx, super_idx, self->_am.get()))
      return false;

   const Molecule &sub_mol = static_cast<const Molecule&>(subgraph);
   const Molecule &super_mol = static_cast<const Molecule&>(supergraph);

   int sub_change = self->_context._query->findReactingCenter(&sub_mol, sub_idx);
   int super_change = self->_context._target.findReactingCenter(&super_mol, super_idx);

   if (super_change == RC_UNMARKED)
      return true;

   // super_change == (RC_UNCHANGED + RC_ORDER_CHANGED) is for changed aromatics
   if (sub_change == RC_NOT_CENTER || sub_change == RC_UNCHANGED)
      return super_change == 0 || super_change == RC_UNCHANGED || super_change == (RC_UNCHANGED + RC_ORDER_CHANGED);

   if (sub_change == RC_CENTER)
      return super_change != 0 && super_change != RC_UNCHANGED && super_change != RC_NOT_CENTER;

   if ((sub_change & super_change) != sub_change)
      return false;

   return true;
}

void ReactionSubstructureMatcher::_Matcher::_removeAtom (const Graph &subgraph, int sub_idx, void *userdata)
{
   ReactionSubstructureMatcher::_Matcher *self = (ReactionSubstructureMatcher::_Matcher *)userdata;

   MoleculeSubstructureMatcher::removeAtom(subgraph, sub_idx, self->_am.get());
}

void ReactionSubstructureMatcher::_Matcher::_addAtom (const Graph &subgraph, const Graph &supergraph,
                           int sub_idx, int super_idx, void *userdata)
{
   ReactionSubstructureMatcher::_Matcher *self = (ReactionSubstructureMatcher::_Matcher *)userdata;

   MoleculeSubstructureMatcher::addAtom(subgraph, supergraph, sub_idx, super_idx, self->_am.get());
}

void ReactionSubstructureMatcher::_Matcher::_addBond (const Graph &subgraph, const Graph &supergraph,
                           int sub_idx, int super_idx, void *userdata)
{
   ReactionSubstructureMatcher::_Matcher *self = (ReactionSubstructureMatcher::_Matcher *)userdata;

   MoleculeSubstructureMatcher::addBond(subgraph, supergraph, sub_idx, super_idx, self->_am.get());
}

